#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from datetime import datetime
from mmcv.cnn import ConvModule, Scale, bias_init_with_prob, normal_init
from mmcv.runner import force_fp32

from mmdet.core import (
    anchor_inside_flags,
    build_anchor_generator,
    bbox2distance,
    bbox_overlaps,
    build_assigner,
    build_sampler,
    distance2bbox,
    images_to_levels,
    multi_apply,
    multiclass_nms,
    reduce_mean,
    unmap,
)
from ..builder import HEADS, build_loss
from .anchor_head import AnchorHead
from ...core.bbox.assigners import print_num_anchor, print_bbox_in_img
from .gfl_head_slim_annotation import GFLSinOut_slim_norm_bbox, DenseBlock
from ..utils.wxz_utils import ste_round_func


def compare_pre_cur_target(
    previous_target, current_target, img_metas, previous_gt_num, current_gt_num
):
    pre_unique = previous_target.unique(dim=0)
    cur_unique = current_target.unique(dim=0)

    written_msgs = ""
    # if pre_unique.equal(cur_unique):
    #     ...
    #     num_gts = pre_unique.size(0)
    #     num_samples_pre = [0 for _ in range(num_gts)]
    #     num_samples_cur = [0 for _ in range(num_gts)]
    #     area = [0 for _ in range(num_gts)]
    #     for gt_index in range(num_gts):
    #         num_samples_pre[gt_index] = (
    #             pre_unique[gt_index] == previous_target
    #         ).sum().cpu().numpy().tolist() // 4

    #         num_samples_cur[gt_index] = (
    #             pre_unique[gt_index] == current_target
    #         ).sum().cpu().numpy().tolist() // 4

    #         area[gt_index] = (
    #             (
    #                 (pre_unique[gt_index][3] - pre_unique[gt_index][1])
    #                 * (pre_unique[gt_index][2] - pre_unique[gt_index][0])
    #             )
    #             .cpu()
    #             .numpy()
    #             .tolist()
    #         )

    #     written_msgs += img_metas[0]["filename"] + "\n"
    #     written_msgs += "gt: " + pre_unique.__str__() + "\n"
    #     written_msgs += f"area: {area}\n"
    #     written_msgs += f"previous: {num_samples_pre}; current: {num_samples_cur}\n\n"

    # else:
    #     pre_len = pre_unique.size(0)
    #     cur_len = pre_unique.size(0)
    #
    #     written_msgs += "miss some gts------------------\n\n"

    pre_len = pre_unique.size(0)
    cur_len = cur_unique.size(0)

    previous_gt_num += pre_len
    current_gt_num += cur_len
    print(img_metas[0]["filename"])
    print(f"previous: {previous_gt_num}; cur: {current_gt_num}")

    return previous_gt_num, current_gt_num

    #  print(written_msgs)
    #  with open("tmp_save/s4ds_anchor_num.log", "a") as f:
    #      f.write(written_msgs)


def print_gt_pred(gt_bboxes, pred_bboxes, pred_bboxes_noshift, img_meta, saved_path):
    """print_gt_pred. Print gt, pred_bboxes, pred_bboxes_noshift in image.

    Args:
        gt_bboxes: Ground Truth bboxes.
        pred_bboxes: Predicted bboxes.
        pred_bboxes_noshift: Predicted bboxes without shifting.
        img_meta: Image basic infomation.
    """

    import cv2

    img_raw = cv2.imread(img_meta["filename"], cv2.IMREAD_COLOR)

    # In cv2, use BGR as default format.
    for i in range(len(gt_bboxes)):
        b = gt_bboxes[i, :]
        # Red
        cv2.rectangle(
            img_raw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (0, 0, 255), 1
        )
    for i in range(len(pred_bboxes)):
        b = pred_bboxes[i, :]
        # Green
        cv2.rectangle(
            img_raw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (0, 255, 0), 1
        )
    for i in range(len(pred_bboxes_noshift)):
        b = pred_bboxes_noshift[i, :]
        # Blue
        cv2.rectangle(
            img_raw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (255, 0, 0), 1
        )
    cv2.imwrite(
        os.path.join(saved_path, img_meta["ori_filename"].split("/")[-1]), img_raw
    )


def print_gt_pred_helper(
    pos_bbox_targets,
    is_in_gts,
    is_in_gts_noshift,
    pos_anchor_centers,
    pos_anchor_centers_noshift,
    bbox_pred,
    img_metas,
    saved_path="tmp_save/",
):
    """print_gt_pred_helper.

    all the args should be based in origin map size.

    Args:
        pos_bbox_targets:
        is_in_gts:
        is_in_gts_noshift:
        pos_anchor_centers:
        pos_anchor_centers_noshift:
        bbox_pred:
        img_metas: assume there is only one image in a batch.
    """
    different_pair = is_in_gts != is_in_gts_noshift

    diff_gt_bboxes = pos_bbox_targets[different_pair]
    diff_center = pos_anchor_centers[different_pair]
    diff_center_noshift = pos_anchor_centers_noshift[different_pair]
    diff_bbox_pred = bbox_pred[different_pair]

    pred_bboxes = distance2bbox(
        diff_center, diff_bbox_pred, max_shape=img_metas[0]["img_shape"]
    )
    pred_bboxes_noshift = distance2bbox(
        diff_center_noshift, diff_bbox_pred, max_shape=img_metas[0]["img_shape"]
    )

    print_gt_pred(
        diff_gt_bboxes, pred_bboxes, pred_bboxes_noshift, img_metas[0], saved_path
    )
